import copy
from typing import Optional, List
import math

import torch
import torch.nn.functional as F
from torch import nn, Tensor
from torch.nn.init import xavier_uniform_, constant_, uniform_, normal_

from util.misc import inverse_sigmoid
from .ops.modules import MSDeformAttn
from .utils import _get_clones, _get_activation_fn



class DeformableTransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.aux_heads = False
        self.class_embed = None
        self.bbox_embed = None
        self.verb_class_embed = None
        self.sub_bbox_embed = None


    def forward(self, src, spatial_shapes, level_start_index, valid_ratios, 
                pos=None, padding_mask=None, topk_inds=None, output_proposals=None, sparse_token_nums=None):
        if self.aux_heads:
            assert output_proposals is not None
        else:
            assert output_proposals is None
            
        output = src
        sparsified_keys = False if topk_inds is None else True
        reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)
        reference_points_orig = reference_points
        pos_orig = pos
        output_proposals_orig = output_proposals
        sampling_locations_all = []
        attn_weights_all = []
        if self.aux_heads:
            enc_inter_outputs_class = []
            enc_inter_outputs_coords = []
            enc_outputs_verb_class = []
            enc_outputs_sub_coord_unact = []
            
#######################################  new  #################################################      
        if sparsified_keys:
            assert topk_inds is not None
            B_, N_, S_, P_ = reference_points.shape
            reference_points = torch.gather(reference_points.view(B_, N_, -1), 1, topk_inds.unsqueeze(-1).repeat(1, 1, S_*P_)).view(B_, -1, S_, P_)
            tgt = torch.gather(output, 1, topk_inds.unsqueeze(-1).repeat(1, 1, output.size(-1)))
            pos = torch.gather(pos, 1, topk_inds.unsqueeze(-1).repeat(1, 1, pos.size(-1)))
            if output_proposals is not None:
                output_proposals = output_proposals.gather(1, topk_inds.unsqueeze(-1).repeat(1, 1, output_proposals.size(-1)))
        else:
            tgt = None
#######################################  new  #################################################
        for lid, layer in enumerate(self.layers):
            # if tgt is None: self-attention / if tgt is not None: cross-attention w.r.t. the target queries
            tgt, sampling_locations, attn_weights = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask, 
                        tgt=tgt if sparsified_keys else None)
            sampling_locations_all.append(sampling_locations)
            attn_weights_all.append(attn_weights)
            if sparsified_keys:                
                if sparse_token_nums is None:
                    output = output.scatter(1, topk_inds.unsqueeze(-1).repeat(1, 1, tgt.size(-1)), tgt)
                else:
                    outputs = []
                    for i in range(topk_inds.shape[0]):
                        outputs.append(output[i].scatter(0, topk_inds[i][:sparse_token_nums[i]].unsqueeze(-1).repeat(1, tgt.size(-1)), tgt[i][:sparse_token_nums[i]]))
                    output = torch.stack(outputs)
            else:
                output = tgt
            
            if self.aux_heads and lid < self.num_layers - 1:
                # feed outputs to aux. heads 
                if self.class_embed is not None:
                    output_class = self.class_embed[lid](tgt)
                    enc_inter_outputs_class.append(output_class)
                
                if self.verb_class_embed is not None:
                    output_verb_class = self.verb_class_embed[lid](tgt)
                    enc_outputs_verb_class.append(output_verb_class)
                
                if self.bbox_embed is not None:
                    output_offset = self.bbox_embed[lid](tgt)
                    output_coords_unact = output_proposals + output_offset
                # values to be used for loss compuation
                enc_inter_outputs_coords.append(output_coords_unact.sigmoid())

                if self.sub_bbox_embed is not None:
                    output_sub_offset = self.sub_bbox_embed[lid](tgt)
                    output_sub_coords_unact = output_proposals + output_sub_offset
                    # values to be used for loss compuation
                    enc_outputs_sub_coord_unact.append(output_sub_coords_unact.sigmoid())

        # Change dimension from [num_layer, batch_size, ...] to [batch_size, num_layer, ...]
        sampling_locations_all = torch.stack(sampling_locations_all, dim=1)
        attn_weights_all = torch.stack(attn_weights_all, dim=1)
        
        out = [output, sampling_locations_all, attn_weights_all]

        if self.aux_heads:
            if self.class_embed is not None:
                out += [enc_inter_outputs_class, enc_inter_outputs_coords]
            if self.sub_bbox_embed is not None:
                out += [enc_outputs_verb_class, enc_outputs_sub_coord_unact]
        
        return out

    @staticmethod
    def get_reference_points(spatial_shapes, valid_ratios, device):
        reference_points_list = []
        for lvl, (H_, W_) in enumerate(spatial_shapes):
            ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
                                          torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
            ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
            ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
            # out-of-reference points have relative coords. larger than 1
            ref = torch.stack((ref_x, ref_y), -1)
            reference_points_list.append(ref)
        reference_points = torch.cat(reference_points_list, 1)
        reference_points = reference_points[:, :, None] * valid_ratios[:, None]
        return reference_points


class DeformableTransformerEncoderLayer(nn.Module):
    def __init__(self,
                 d_model=256, d_ffn=1024,
                 dropout=0.1, activation="relu",
                 n_levels=4, n_heads=8, n_points=4):
        super().__init__()
        # self attention
        self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)     
        self.dropout1 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(d_model)
        
        self.build_ffn(dropout, activation, d_model, d_ffn)

    @staticmethod
    def with_pos_embed(tensor, pos):
        return tensor if pos is None else tensor + pos

    def build_ffn(self, dropout, activation, d_model, dim_feedforward):
        # Implementation of Feedforward model  
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.activation = _get_activation_fn(activation)

    def forward_ffn(self, src):
        src2 = self.linear2(self.dropout2(self.activation(self.linear1(src))))
        src = src + self.dropout3(src2)
        src = self.norm2(src)
        return src

    def forward(self, src, pos, reference_points, spatial_shapes, level_start_index, padding_mask=None, tgt=None):
        if tgt is None:
            # self attention N, 35x35, 256 
            q = self.with_pos_embed(src, pos)
            k = reference_points    #change
            v = src
            src2, sampling_locations, attn_weights = self.self_attn(q, k, v, 
                                                    spatial_shapes, level_start_index, padding_mask)
            src = src + self.dropout1(src2)
            src = self.norm1(src)
            # ffn
            src = self.forward_ffn(src)
            return src, sampling_locations, attn_weights
        else:
            q = self.with_pos_embed(tgt, pos)
            k = reference_points    #change
            v = src
            # self attention
            tgt2, sampling_locations, attn_weights = self.self_attn(q, k, v, spatial_shapes,
                                                                    level_start_index, padding_mask)
            tgt = tgt + self.dropout1(tgt2)
            tgt = self.norm1(tgt)

            # ffn
            tgt = self.forward_ffn(tgt)

            return tgt, sampling_locations, attn_weights
